import numpy as np
import yaml
import datetime as dt
from datetime import datetime
from util.process_lib import readloc, match_coordinates, SEC2, data_interpolate, bandpass_filter, generate_slices
from util.gpr_io_dt1 import readdt1, readdt1Header
from typing import Dict

class PipeCommunication:
    """TransmissionClient 专用管道通信模块"""
    def __init__(self, pipe_name=r'\\.\pipe\transmission_pipe'):
        self.pipe_name = pipe_name
        self.pipe_handle = None
        self.running = False
        self.listener_thread = None
        self.transmission_client = None  # 关联TransmissionClient实例

    def set_client(self, client):
        self.transmission_client = client

    def connect(self):
        """连接到C++的命名管道"""
        try:
            self.pipe_handle = win32file.CreateFile(
                self.pipe_name,
                win32file.GENERIC_READ | win32file.GENERIC_WRITE,
                0, None, win32file.OPEN_EXISTING, 0, None
            )
            print(f"[Transmission] 管道连接成功: {self.pipe_name}")
            self.running = True
            self.listener_thread = threading.Thread(target=self._listen, daemon=True)
            self.listener_thread.start()
            return True
        except Exception as e:
            print(f"[Transmission] 管道连接失败: {str(e)}")
            return False

    def send(self, data: Dict):
        """向C++发送消息"""
        if not self.pipe_handle or not self.running:
            print("[Transmission] 管道未连接，发送失败")
            return False
        try:
            msg = json.dumps(data) + "\r\n"
            win32file.WriteFile(self.pipe_handle, msg.encode('utf-8'))
            print(f"[Transmission] 发送到C++: {data}")
            return True
        except Exception as e:
            print(f"[Transmission] 发送失败: {str(e)}")
            return False

    def _listen(self):
        """监听C++指令"""
        while self.running:
            try:
                _, data = win32file.ReadFile(self.pipe_handle, 4096)
                if data:
                    msg = data.decode('utf-8').strip()
                    print(f"[Transmission] 收到C++指令: {msg}")
                    self._handle_command(msg)
            except Exception as e:
                if "断开的管道" in str(e):
                    print("[Transmission] 管道已断开")
                    self.running = False
                    break
                print(f"[Transmission] 监听异常: {str(e)}")
            threading.Event().wait(0.1)

    def _handle_command(self, msg_str: str):
        """处理C++指令"""
        if not self.transmission_client:
            self.send({"type": "error", "msg": "未初始化TransmissionClient"})
            return

        try:
            msg = json.loads(msg_str)
            cmd_type = msg.get("type")

            # 执行数据评估
            if cmd_type == "evaluate":
                self.transmission_client.evaluate()
                self.send({
                    "type": "evaluate_done",
                    "status": "success",
                    "has_3d": self.transmission_client.volume_data is not None,
                    "has_slices": self.transmission_client.slices_data is not None
                })

            # 获取3D数据
            elif cmd_type == "get_3d":
                if self.transmission_client.volume_data:
                    self.send({
                        "type": "3d_data",
                        "data": self.transmission_client.volume_data
                    })
                else:
                    self.send({"type": "error", "msg": "3D数据未生成"})

            # 获取切片数据
            elif cmd_type == "get_slices":
                if self.transmission_client.slices_data:
                    self.send({
                        "type": "slices_data",
                        "count": len(self.transmission_client.slices_data),
                        "data": self.transmission_client.slices_data
                    })
                else:
                    self.send({"type": "error", "msg": "切片数据未生成"})

            # 退出指令
            elif cmd_type == "exit":
                self.running = False
                self.transmission_client.close()
                sys.exit(0)

            else:
                self.send({"type": "error", "msg": f"未知指令: {cmd_type}"})

        except json.JSONDecodeError:
            self.send({"type": "error", "msg": "无效的JSON格式"})
        except Exception as e:
            self.send({"type": "error", "msg": str(e)})

    def close(self):
        """关闭管道"""
        self.running = False
        if self.listener_thread and self.listener_thread.is_alive():
            self.listener_thread.join(1.0)
        if self.pipe_handle:
            try:
                win32file.CloseHandle(self.pipe_handle)
            except Exception as e:
                print(f"[Transmission] 关闭管道失败: {str(e)}")


class TransmissionClient:
    def __init__(self, config_file):
        self.connected = False

        self.config_file = config_file
        self.start_time = None
        self.dt = None
        self.seatrac_path = None
        self.guidance_path = None
        self.dt_path = None
        self.location_path = None
        self.do_filter = None
        self.filter_low = None
        self.filter_high = None
        self.do_gain = None
        self.attenuation = None
        self.start_gain = None
        self.maximum_gain = None
        self.do_interpolate = None
        self.d_horizontal = None
        self.dz = None
        self.vel = None
        self.do_slice = None
        self.d_horizontal_slice = None
        self.neighborhood_radius = None
        self.slice_interval = None
        self.slice_thickness = None
        self.output_trace_location_file = None
        self.output_data_file = None

        self.timestamps = None
        self.coordinates = None
        self.info = None
        self.N_traces = None
        self.N_traces = None
        self.npts = None
        self.sample_rate = None
        self.data_times = None
        self.matched_coordinates = None
        self.data_raw = None
        self.data = None

        self.volume_data = None
        self.slices_data = None

        self.connect()

    def connect(self):
        config = self._load_config()
        self.config = config["param"]

        self.start_time = datetime.strptime(self.config["start_time"],"%Y-%m-%d %H:%M:%S.%f")
        self.dt = self.config["dt"]
        self.seatrac_path = self.config["seatrac_path"]
        self.guidance_path = self.config["guidance_path"]
        self.dt_path = self.config["dt_path"]
        self.location_path = self.config["location_path"]
        self.do_filter = self.config["do_filter"]
        self.filter_low = self.config["filter_low"]
        self.filter_high = self.config["filter_high"]
        self.do_gain = self.config["do_gain"]
        self.attenuation = self.config["attenuation"]
        self.start_gain = self.config["start_gain"]
        self.maximum_gain = self.config["maximum_gain"]
        self.do_interpolate = self.config["do_interpolate"]
        self.d_horizontal = self.config["d_horizontal"]
        self.dz = self.config["dz"]
        self.vel = self.config["vel"]
        self.do_slice = self.config["do_slice"]
        self.d_horizontal_slice = self.config["d_horizontal_slice"]
        self.neighborhood_radius = self.config["neighborhood_radius"]
        self.slice_interval = self.config["slice_interval"]
        self.slice_thickness = self.config["slice_thickness"]
        self.output_trace_location_file = self.config["output_trace_location_file"]
        self.output_data_file = self.config["output_data_file"]

        # 读取坐标数据
        self.timestamps, self.coordinates = readloc(self.location_path)

        # 读取雷达道头数据
        self.info = readdt1Header(self.dt_path.replace('.dt1', '.hd'))
        self.N_traces = self.info["N_traces"]
        self.npts = self.info['N_pts_per_trace']
        self.sample_rate = self.info['Total_time_window'] / self.npts

        # 生成数据时间列表
        self.data_times = [self.start_time + dt.timedelta(seconds=i * self.dt) for i in range(self.N_traces)]

        # 匹配坐标数据
        self.matched_coordinates = match_coordinates(self.timestamps, self.coordinates, self.data_times)
        # 输出雷达数据每道的坐标
        with open(self.output_trace_location_file, 'w') as f:
            f.write("Trace,Longitude,Latitude,Depth\n")
            for i, coord in enumerate(self.matched_coordinates):
                f.write(f"{i+1},{coord[0]},{coord[1]},{coord[2]}\n")

        # 读取雷达数据
        self.data_raw = readdt1(self.dt_path)

        self.connected = True

    def _load_config(self) -> Dict:
        """加载配置文件（从yaml读取）"""
        try:
            with open(self.config_file, "r", encoding="utf-8") as f:
                self.config = yaml.safe_load(f)  # 读取配置到实例属性
            return self.config  # 正确返回实例属性，而非未定义的变量
        except Exception as e:
            raise ValueError(f"加载配置文件失败：{str(e)}")  # 明确报错原因

    def evaluate(self):
        """主函数：仅负责流程调度，调用各功能子函数"""
        print("=== 开始数据评估与处理 ===")

        # 1. 估算网格大小（独立子函数）
        self._estimate_grid_size()

        # 2. 数据预处理（滤波+增益，独立子函数）
        self._process_raw_data()

        # 3. 3D插值处理（独立子函数）
        self._process_3d_interpolation()

        # 4. 2D切片处理（独立子函数）
        self._process_2d_slices()

        print("=== 数据评估与处理完成 ===")

    def _estimate_grid_size(self):
        """子函数1：仅负责估算三维网格和二维切片大小（私有方法，内部调用）"""
        print("📊 估算网格大小...")
        try:
            # 提取经纬度范围（复用原逻辑，仅做格式整理）
            valid_coords = [coord for coord in self.matched_coordinates if
                            coord[0] is not None and coord[1] is not None]
            if not valid_coords:
                raise ValueError("无有效经纬度坐标，无法估算网格大小")

            lons = [coord[0] for coord in valid_coords]
            lats = [coord[1] for coord in valid_coords]
            min_lon, max_lon = min(lons), max(lons)
            min_lat, max_lat = min(lats), max(lats)
            lon_range = max_lon - min_lon
            lat_range = max_lat - min_lat

            # 计算最大深度
            max_depth = self.sample_rate * self.npts * self.vel / 2

            # 估算3D网格大小（仅当开启插值时）
            if self.do_interpolate:
                grid_x = int(lon_range / self.d_horizontal) + 1
                grid_y = int(lat_range / self.d_horizontal) + 1
                grid_z = int(max_depth / self.dz) + 1
                total_voxels = grid_x * grid_y * grid_z
                print(f"✅ 三维网格估算结果: {grid_x} x {grid_y} x {grid_z} = {total_voxels} 个体素")
                # 可选：将结果存为实例属性，供后续使用
                self.estimated_3d_grid = (grid_x, grid_y, grid_z)

            # 估算2D切片大小（仅当开启切片时）
            if self.do_slice:
                slice_x = int(lon_range / self.d_horizontal_slice) + 1
                slice_y = int(lat_range / self.d_horizontal_slice) + 1
                num_slices = (self.npts - self.slice_thickness) // self.slice_interval
                total_pixels_per_slice = slice_x * slice_y
                print(
                    f"✅ 二维切片估算结果: 单切片 {slice_x} x {slice_y} = {total_pixels_per_slice} 像素，共 {num_slices} 个切片")
                # 可选：将结果存为实例属性
                self.estimated_2d_slice = (slice_x, slice_y, num_slices)

        except Exception as e:
            print(f"❌ 网格大小估算失败: {str(e)}")
            # 失败时初始化空属性，避免后续调用报错
            self.estimated_3d_grid = None
            self.estimated_2d_slice = None

    def _process_raw_data(self):
        """子函数2：仅负责原始数据预处理（滤波+增益，私有方法）"""
        print("🔧 开始数据预处理...")
        # 初始化处理后的数据为原始数据
        self.data = self.data_raw.copy()  # 用copy避免修改原始数据

        # 1. 带通滤波（仅当开启时）
        if self.do_filter:
            print("   - 执行带通滤波...")
            try:
                # 单位转换（原逻辑保留，仅优化变量名增强可读性）
                sample_rate_hz = 1e9 / self.fsample_rate  # 纳秒 → Hz
                filter_low_hz = self.filter_low * 1e6  # MHz → Hz
                filter_high_hz = self.filter_high * 1e6  # MHz → Hz

                self.data = bandpass_filter(
                    data=self.data,
                    lowcut=filter_low_hz,
                    highcut=filter_high_hz,
                    fs=sample_rate_hz
                )
                print("   ✅ 带通滤波完成")
            except Exception as e:
                print(f"   ❌ 带通滤波失败: {str(e)}，将使用原始数据")
                self.data = self.data_raw.copy()

        # 2. 增益处理（仅当开启时）
        if self.do_gain:
            print("   - 执行增益处理...")
            try:
                self.data = SEC2(
                    data=self.data,
                    sample_rate=self.sample_rate,
                    vel=self.vel,
                    attenuation=self.attenuation,
                    start_gain=self.start_gain,
                    maximum_gain=self.maximum_gain
                )
                print("   ✅ 增益处理完成")
            except Exception as e:
                print(f"   ❌ 增益处理失败: {str(e)}，将使用滤波后（或原始）数据")

    def _process_3d_interpolation(self):
        """子函数3：仅负责3D插值处理（含格式转换为Plotly友好格式，私有方法）"""
        self.volume_data = None  # 初始化空属性
        if not self.do_interpolate:
            print("📌 未开启3D插值，跳过该步骤")
            return

        print("🔍 开始3D插值处理...")
        try:
            # 调用插值工具函数（原逻辑保留）
            grid_x_mesh, grid_y_mesh, grid_z_mesh, grid_values = data_interpolate(
                self.data,self.matched_coordinates,self.d_horizontal, self.dz,self.vel,self.sample_rate)

            # 格式转换（展平为1D列表，适配Plotly）
            x = grid_x_mesh.flatten().tolist()
            y = grid_y_mesh.flatten().tolist()
            z = grid_z_mesh.flatten().tolist()
            values = grid_values.flatten().tolist()

            # 构建最终的3D体积数据结构
            self.volume_data = {
                "x": x,
                "y": y,
                "z": z,
                "value": values,
                "grid_shape": grid_x_mesh.shape  # 可选：保留网格原始形状，方便后续反推
            }
            print(f"✅ 3D插值完成，共 {len(values)} 个体素数据")

        except Exception as e:
            print(f"❌ 3D插值处理失败: {str(e)}")
            self.volume_data = None

    def _process_2d_slices(self):
        """子函数4：仅负责2D切片处理（含格式转换为Plotly友好格式，私有方法）"""
        self.slices_data = None  # 初始化空属性
        if not self.do_slice:
            print("📌 未开启2D切片，跳过该步骤")
            return

        print("📄 开始2D切片处理...")
        try:
            # 调用切片生成工具函数（原逻辑保留，修正拼写错误：silce_thickness → slice_thickness）
            lon_min, lon_max, lat_min, lat_max, grid_values_list, depths = generate_slices(
                self.data, self.matched_coordinates, self.d_horizontal_slice, self.vel, self.sample_rate,
                self.slice_interval, self.slice_thickness, self.neighborhood_radius
            )

            # 格式转换：处理depths（确保为列表）
            depths = depths.tolist() if isinstance(depths, np.ndarray) else depths
            if len(depths) != len(grid_values_list):
                raise ValueError(f"切片数量与深度数量不匹配（切片{len(grid_values_list)}个，深度{len(depths)}个）")

            # 构建切片列表（适配Plotly）
            self.slices_data = []
            for idx, (grid, depth) in enumerate(zip(grid_values_list, depths)):
                # 转置网格（符合Plotly的x/y轴习惯）
                grid_array = np.array(grid) if not isinstance(grid, np.ndarray) else grid
                grid_transposed = grid_array.T.tolist()

                # 生成切片的x/y坐标列表（对应经纬度）
                slice_x = [lon_min + j * self.d_horizontal_slice for j in range(len(grid_transposed[0]))]
                slice_y = [lat_min + i * self.d_horizontal_slice for i in range(len(grid_transposed))]

                # 添加单个切片数据
                self.slices_data.append({
                    "slice_index": idx,
                    "depth": depth,
                    "x": slice_x,
                    "y": slice_y,
                    "z": grid_transposed  # Plotly热力图/等高线图需要的z值矩阵
                })

            print(f"✅ 2D切片处理完成，共生成 {len(self.slices_data)} 个切片")

        except Exception as e:
            print(f"❌ 2D切片处理失败: {str(e)}")
            self.slices_data = None

    def close(self):
        self.start_time = None
        self.dt = None
        self.seatrac_path = None
        self.guidance_path = None
        self.dt_path = None
        self.location_path = None
        self.do_filter = None
        self.filter_low = None
        self.filter_high = None
        self.do_gain = None
        self.attenuation = None
        self.start_gain = None
        self.maximum_gain = None
        self.do_interpolate = None
        self.d_horizontal = None
        self.dz = None
        self.vel = None
        self.do_slice = None
        self.d_horizontal_slice = None
        self.neighborhood_radius = None
        self.slice_interval = None
        self.slice_thickness = None
        self.output_trace_location_file = None
        self.output_data_file = None

        self.timestamps = None
        self.coordinates = None
        self.info = None
        self.N_traces = None
        self.N_traces = None
        self.npts = None
        self.sample_rate = None
        self.data_times = None
        self.matched_coordinates = None
        self.data_raw = None
        self.data = None

        self.volume_data = None
        self.slices_data = None
        print(f"🔌 任务客户端已关闭")

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("请传入配置文件路径：python transmission.py ./configs/transmission.yaml")
        sys.exit(1)
    config_path = sys.argv[1]
    client = TransmissionClient(config_path)
    # 保持运行以接收管道指令
    try:
        while client.pipe.running:
            threading.Event().wait(1)
    except KeyboardInterrupt:
        print("用户中断，退出程序")
    finally:
        client.close()

